import json
import pandas as pd
import numpy as np
import lightgbm as lgb
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error


def load_data(filepath):
    # 加载 JSON 数据
    with open(filepath, 'r') as file:
        data = json.load(file)

    # 转换为 DataFrame
    df = pd.DataFrame(data)
    return df


def train_and_predict(boosting_type, learning_rate, feature_fraction, bagging_fraction, data_file='../data/test2.json',
                             forecast_file='../data/test3.json'):
    # 加载数据
    df = load_data(data_file)

    # 将数据列转换为 numpy 数组
    X = df['x'].values.reshape(-1, 1)
    y = df['y'].values

    # 划分训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    # 创建 LightGBM 数据矩阵
    train_data = lgb.Dataset(X_train, label=y_train)

    # 设置模型参数
    params = {
        'boosting_type': boosting_type,
        'objective': 'regression',
        'metric': 'rmse',
        'num_leaves': 31,
        'learning_rate': learning_rate,
        'feature_fraction': feature_fraction,
        'bagging_fraction': bagging_fraction,
        'bagging_freq': 5,
        'verbose': 0
    }

    # 训练模型
    model = lgb.train(params, train_data, num_boost_round=100)

    # 预测未来的数据点
    last_x = df['x'].iloc[-1]  # 获取最后一个 x 值
    future_x = [last_x + 0.1 * i for i in range(1, 11)]
    future_x = np.array(future_x).reshape(-1, 1)
    predictions = model.predict(future_x)

    # 创建预测数据的 JSON 对象
    last_x = df.index[-1] + 0.1
    forecast_data = [{'x': round(last_x * 0.1 + i * 0.1, 1), 'y': int(round(y))} for i, y in enumerate(predictions)]

    # # 写入结果到 JSON 文件
    # results = [{'x': float(fx), 'y': float(pred)} for fx, pred in zip(future_x.flatten(), predictions)]
    with open(forecast_file, 'w') as outfile:
        json.dump(forecast_data, outfile)

    print("完成预测，结果已写入到 test3.json")


# if __name__ == '__main__':
#     # 示例参数调用
#     train_and_predict('gbdt', 0.05, 0.9, 0.8)
